package com.datastructures2.graph.最小生成树;

import edu.princeton.cs.algs4.*;

import java.util.Arrays;

/**
 * @author MaoLin Wang
 * @date 2020/2/2314:39
 */
public class PrimMST {
    private static final double FLOATING_POINT_EPSILON = 1E-12;
    /**
     * edgeTo[v] 顶点v距离树最近的边
     */
    private Edge[] edgeTo;

    /**
     * distTo[v] 保存顶点v到树顶点最小边的权重
     */
    private double[] distTo;

    private boolean[] marked;

    /**
     * 最小生成树的权重
     */
    private double weight=0.0;

    /**
     * 有效的横切边：索引优先队列，顶点v作为索引，disTo[v]作为索引关联的值
     */
    private IndexMinPQ<Double>pq;

    public PrimMST(EdgeWeightedGraph G){
        edgeTo=new Edge[G.V()];
        distTo=new double[G.V()];
        marked=new boolean[G.V()];

        for (int v = 0; v < G.V(); v++) {
            //初始化为无穷大
            distTo[v]=Double.POSITIVE_INFINITY;
        }
        pq=new IndexMinPQ<>(G.V());

        //如果是连通的，则直接走prim()，这个循环不需要加
        for (int v = 0; v < G.V(); v++) {
            if (!marked[v]){
                prim(G,v);
            }
        }
         check(G);
    }

    private void prim(EdgeWeightedGraph G,int v) {
        distTo[v]=0.0;
        //将当前顶点作为索引，当前顶点距离树权重最小的边的权重作为关联的key加入索引优先队列
        pq.insert(v,distTo[v]);

        while (!pq.isEmpty()){
            //取出权重最小的索引（顶点）
            int w = pq.delMin();
            visit(G,w);
        }
    }

    private void visit(EdgeWeightedGraph G, int v) {
        marked[v]=true;
        weight+=distTo[v];
        for (Edge e:G.adj(v)){
            int w=e.other(v);

            if (marked[w]){
                continue;
            }
            if (e.weight()<distTo[w]){//只有当新的边的权重小于之前的distTo[w]值的时候，才更新值
                edgeTo[w]=e;
                distTo[w]=e.weight();
                if (pq.contains(w)){ //索引优先队列有这个索引，更新对应的权重值
                    pq.changeKey(w,distTo[w]);
                }else {//否则插入新的
                    pq.insert(w,distTo[w]);
                }
            }
        }
    }
    ///check()
    private boolean check(EdgeWeightedGraph G) {

        // check weight
        double totalWeight = 0.0;
        for (Edge e : edges()) {
            totalWeight += e.weight();
        }
        if (Math.abs(totalWeight - weight()) > FLOATING_POINT_EPSILON) {
            System.err.printf("Weight of edges does not equal weight(): %f vs. %f\n", totalWeight, weight());
            return false;
        }

        // check that it is acyclic
        UF uf = new UF(G.V());
        for (Edge e : edges()) {
            int v = e.either(), w = e.other(v);
            if (uf.connected(v, w)) {
                System.err.println("Not a forest");
                return false;
            }
            uf.union(v, w);
        }

        // check that it is a spanning forest
        for (Edge e : G.edges()) {
            int v = e.either(), w = e.other(v);
            if (!uf.connected(v, w)) {
                System.err.println("Not a spanning forest");
                return false;
            }
        }

        // check that it is a minimal spanning forest (cut optimality conditions)
        for (Edge e : edges()) {

            // all edges in MST except e
            uf = new UF(G.V());
            for (Edge f : edges()) {
                int x = f.either(), y = f.other(x);
                if (f != e) uf.union(x, y);
            }

            // check that e is min weight edge in crossing cut
            for (Edge f : G.edges()) {
                int x = f.either(), y = f.other(x);
                if (!uf.connected(x, y)) {
                    if (f.weight() < e.weight()) {
                        System.err.println("Edge " + f + " violates cut optimality conditions");
                        return false;
                    }
                }
            }

        }

        return true;
    }
    public Iterable<Edge> edges(){
        Queue<Edge> mst = new Queue<Edge>();
        for (int v = 0; v < edgeTo.length; v++) {
            Edge e = edgeTo[v];
            if (e != null) {
                mst.enqueue(e);
            }
        }
        return mst;
    }
    public double weight(){
        return weight;
    }

    public static void main(String[] args) {
        In in = new In("D:\\JavaProject\\算法\\数据结构和算法\\DATAStructures\\alg\\src\\main\\java\\com\\datastructures2\\graph\\最小生成树\\tinyEWG.txt");
        EdgeWeightedGraph G = new EdgeWeightedGraph(in);
        PrimMST mst = new PrimMST(G);
        for (Edge e : mst.edges()) {
            StdOut.println(e);
        }
        StdOut.printf("%.5f\n", mst.weight());

    }
}
